package cn.jiaqiang.chatgpt.data.domain.openai.service.rule.impl.filter;


import cn.jiaqiang.chatgpt.data.domain.openai.annotation.LogicStrategy;
import cn.jiaqiang.chatgpt.data.domain.openai.model.aggregates.ChatProcessAggregate;
import cn.jiaqiang.chatgpt.data.domain.openai.model.entity.RuleLogicEntity;
import cn.jiaqiang.chatgpt.data.domain.openai.model.entity.UserAccountEntity;
import cn.jiaqiang.chatgpt.data.domain.openai.model.entity.UserModelEntity;
import cn.jiaqiang.chatgpt.data.domain.openai.model.valobj.LogicCheckTypeVO;
import cn.jiaqiang.chatgpt.data.domain.openai.repository.IUserModelRepository;
import cn.jiaqiang.chatgpt.data.domain.openai.service.rule.ILogicFilter;
import cn.jiaqiang.chatgpt.data.domain.openai.service.rule.factory.DefaultLogicFactory;
import cn.jiaqiang.chatgpt.data.types.enums.ChatGPTModel;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;

/**
 * @description 允许访问的模型过滤
 */
@Slf4j
@Component
@LogicStrategy(logicMode = DefaultLogicFactory.LogicModel.MODEL_TYPE)
public class ModelTypeFilter implements ILogicFilter<UserAccountEntity> {

    @Resource
    private IUserModelRepository userModelRepository;

    @Override
    public RuleLogicEntity<ChatProcessAggregate> filter(ChatProcessAggregate chatProcess, UserAccountEntity data) throws Exception {
        String modelType = chatProcess.getModel();
        String userId = chatProcess.getUserId();
        log.info("userId = {} 开始进行【可用模型拦截】, 当前模型: {}", userId, modelType);

        if ((null == data || data.getSurplusQuota() == 0) &&
                ChatGPTModel.GPT_3_5_TURBO.getCode().equals(chatProcess.getModel())) {
            // 当前是免费次数的校验, 直接返回成功
            log.info("当前是免费次数的校验 userId = {}, 【可用模型拦截】结果: {}", userId, LogicCheckTypeVO.SUCCESS.getInfo());
            return RuleLogicEntity.<ChatProcessAggregate>builder()
                    .type(LogicCheckTypeVO.SUCCESS)
                    .data(chatProcess)
                    .build();
        }

        // 1. 通过用户ID 和 模型去查额度 -> 模型和额度都进行了拦截
        UserModelEntity userModelEntity =
                userModelRepository.queryUserModelEntityWithUserIdModel(Long.parseLong(userId), chatProcess.getModel());
        if (userModelEntity != null && userModelEntity.getQuota() != 0) {
            log.info("userId = {}, 【可用模型拦截】结果: {}", userId, LogicCheckTypeVO.SUCCESS.getInfo());
            return RuleLogicEntity.<ChatProcessAggregate>builder()
                    .type(LogicCheckTypeVO.SUCCESS)
                    .data(chatProcess)
                    .build();
        }

        // 3. 模型校验拦截
        log.info("userId = {}, 【可用模型拦截】结果: {} -> 不支持的模型: {}",
                userId, LogicCheckTypeVO.REFUSE.getInfo(), modelType);
        return RuleLogicEntity.<ChatProcessAggregate>builder()
                .type(LogicCheckTypeVO.REFUSE)
                .info("当前账户 " + modelType + " 模型次数不足！可以联系客服升级账户。")
                .data(chatProcess)
                .build();
    }

}
